'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 4_few_shot_prompttemplate.py
* @Time: 2025/10/30
* @All Rights Reserve By Brtc
'''
import dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, FewShotPromptTemplate, FewShotChatMessagePromptTemplate
from langchain_openai import ChatOpenAI

dotenv.load_dotenv()

#1、构建示例模板与示例
examples = [
    {"question": "帮我计算下2+2等于多少？", "answer": "4"},
    {"question": "帮我计算下2+3等于多少？", "answer": "5"},
    {"question": "帮我计算下20*15等于多少？", "answer": "300"},
]
example_prompt = ChatPromptTemplate.from_messages([
    ("human", "{question}"),
    ("ai", "{answer}")
])
#2、构建少量示例提示模板
few_shot_prompt = FewShotChatMessagePromptTemplate(
    example_prompt = example_prompt,
    examples = examples,
)
print("少量提示词模板:", few_shot_prompt.format())
#3、构建最终提示词模板
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个计算非常牛逼的机器人！"),
    few_shot_prompt,
    ("human", "{question}")
])
#4、创建大预言模型
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
chain = prompt|llm|StrOutputParser()
#5、调用链并获取结果
print(chain.invoke("帮我计算以下14*15等于多少？,直接给出答案不要废话"))
